"""
This file is part of the openPMD-api.

Copyright 2021 openPMD contributors
Authors: Axel Huebl
License: LGPLv3+
"""
import math

import numpy as np


def particles_to_dataframe(particle_species, slice=None):
    """
    Load all records of a particle species into a Pandas DataFrame.

    Parameters
    ----------
    particle_species : openpmd_api.ParticleSpecies
        A ParticleSpecies class in openPMD-api.
    slice : np.s_, optional
        A numpy slice that can be used to load only a sub-selection of
        particles.

    Returns
    -------
    pandas.DataFrame
        A pandas dataframe with particles as index and openPMD record
        components of the particle_species as columns.

    Raises
    ------
    ImportError
        Raises an exception if pandas is not installed

    See Also
    --------
    numpy.s_ : the slice object to sub-select
    openpmd_api.BaseRecordComponent.available_chunks : available chunks that
        are optimal arguments for the slice parameter
    pandas.DataFrame : the central dataframe object created here
    """
    # import pandas here for a lazy import
    try:
        import pandas as pd
        found_pandas = True
    except ImportError:
        found_pandas = False

    if not found_pandas:
        raise ImportError("pandas NOT found. Install pandas for DataFrame "
                          "support.")
    if slice is None:
        slice = np.s_[()]

    columns = {}

    for record_name, record in particle_species.items():
        for rc_name, rc in record.items():
            if record.scalar:
                column_name = record_name
            else:
                column_name = record_name + "_" + rc_name
            columns[column_name] = rc[slice]
            particle_species.series_flush()
            if not math.isclose(1.0, rc.unit_SI):
                columns[column_name] = np.multiply(
                    columns[column_name], rc.unit_SI)

    df = pd.DataFrame(columns)

    # set a header for the first column (row index)
    #   note: this is NOT the particle id
    df.index.name = "row"

    return df


def iterations_to_dataframe(series, species_name):
    """
    Load all iterations of a particle species into a Pandas DataFrame.

    Parameters
    ----------
    series : openpmd_api.Series
        A Series class in openPMD-api.
    species_name : string
        The name of a particle species.

    Returns
    -------
    pandas.DataFrame
        A pandas dataframe with particles as index and openPMD record
        components of the particle_species as columns. Particles might be
        repeated over multiple iterations and an "iteration" column is
        added.

    Raises
    ------
    ImportError
        Raises an exception if pandas is not installed

    See Also
    --------
    pandas.DataFrame : the central dataframe object created here
    """
    # import pandas here for a lazy import
    try:
        import pandas as pd
    except ImportError:
        raise ImportError("pandas NOT found. Install pandas for DataFrame "
                          "support.")

    df = pd.concat(
        (
            series.iterations[i]
            .particles[species_name]
            .to_df()
            .assign(iteration=i)
            for i in series.iterations
        ),
        axis=0,
        ignore_index=True,
    )

    return df


def iterations_to_cudf(series, species_name):
    """
    Load all iterations of a particle species into a cuDF DataFrame.

    Parameters
    ----------
    series : openpmd_api.Series
        A Series class in openPMD-api.
    species_name : string
        The name of a particle species.

    Returns
    -------
    cudf.DataFrame
        A cuDF (RAPIDS) dataframe with particles as index and openPMD record
        components of the particle_species as columns. Particles might be
        repeated over multiple iterations and an "iteration" column is
        added.

    Raises
    ------
    ImportError
        Raises an exception if cuDF (RAPIDS) is not installed

    See Also
    --------
    cudf.DataFrame : the central dataframe object created here
    """
    # import pandas here for a lazy import
    try:
        import pandas  # noqa
    except ImportError:
        raise ImportError("pandas NOT found. Install pandas for DataFrame "
                          "support.")
    # import cudf here for a lazy import
    try:
        import cudf
    except ImportError:
        raise ImportError("cudf NOT found. Install RAPIDS for CUDA DataFrame "
                          "support.")

    cdf = cudf.concat(
        (
            cudf.from_pandas(
                series.iterations[i]
                      .particles[species_name]
                      .to_df()
                      .assign(iteration=i)
            )
            for i in series.iterations
        ),
        axis=0,
        ignore_index=True,
    )

    return cdf
